{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Batch: Tianshou's Core Data Structure\n",
    "\n",
    "The `Batch` class is Tianshou's fundamental data structure for efficiently storing and manipulating heterogeneous data in reinforcement learning. This tutorial provides comprehensive guidance on understanding its conceptual foundations, operational behavior, and best practices.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "from typing import cast\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.distributions import Categorical, Normal\n",
    "\n",
    "from tianshou.data import Batch\n",
    "from tianshou.data.types import ActBatchProtocol, ObsBatchProtocol, RolloutBatchProtocol"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Introduction: Why Batch?\n",
    "\n",
    "### The Challenge in Reinforcement Learning\n",
    "\n",
    "Reinforcement learning algorithms face a fundamental data management challenge:\n",
    "\n",
    "1. **Diverse Data Requirements**: Different RL algorithms need different data fields:\n",
    "   - Basic algorithms: `state`, `action`, `reward`, `done`, `next_state`\n",
    "   - Actor-Critic: additionally `advantages`, `returns`, `values`\n",
    "   - Policy Gradient: additionally `log_probs`, `old_log_probs`\n",
    "   - Off-policy: additionally `priority_weights`\n",
    "\n",
    "2. **Heterogeneous Observation Spaces**: Environments return diverse observation types:\n",
    "   - Simple: vectors (`np.array([1.0, 2.0, 3.0])`)\n",
    "   - Complex: images (`np.array(shape=(84, 84, 3))`)\n",
    "   - Hybrid: dictionaries combining multiple modalities\n",
    "   ```python\n",
    "   obs = {\n",
    "       'camera': np.array(shape=(64, 64, 3)),\n",
    "       'velocity': np.array([1.2, 0.5]),\n",
    "       'inventory': np.array([5, 2, 0])\n",
    "   }\n",
    "   ```\n",
    "\n",
    "3. **Data Flow Across Components**: Data must flow seamlessly through:\n",
    "   - Collectors (gathering experience from environments)\n",
    "   - Replay Buffers (storing and sampling transitions)\n",
    "   - Policies and Algorithms (learning and inference)\n",
    "\n",
    "### Why Not Alternatives?\n",
    "\n",
    "#### Plain Dictionaries\n",
    "Dictionaries lack essential features\n",
    "```python\n",
    "data = {'obs': np.array([1, 2]), 'reward': np.array([1.0, 2.0])}\n",
    "```\n",
    "\n",
    "They would work in principle but has no shape/length semantics, no indexing, and no type safety.\n",
    "\n",
    "#### TensorDict\n",
    "While `TensorDict` (used in `pytorch-rl`) is a powerful alternative:\n",
    "- **Batch supports arbitrary objects**, not just tensors (useful for object-dtype arrays, custom types)\n",
    "- **Batch has better type checking** via `BatchProtocol` (enables IDE autocompletion)\n",
    "- **Batch preceded TensorDict** and provides a stable foundation for Tianshou\n",
    "- **TensorDict isn't part of core PyTorch** (external dependency)\n",
    "\n",
    "### What is Batch?\n",
    "\n",
    "**Batch = Dictionary + Array hybrid with RL-specific features**\n",
    "\n",
    "Key capabilities:\n",
    "- **Dict-like**: Key-value storage with attribute access (`batch.obs`, `batch.reward`)\n",
    "- **Array-like**: Shape, indexing, slicing (`batch[0]`, `batch[:10]`, `batch.shape`)\n",
    "- **Hierarchical**: Nested structures for complex data\n",
    "- **Type-safe**: Protocol-based typing for IDE support\n",
    "- **RL-aware**: Special handling for distributions, missing values, heterogeneous aggregation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Core Concepts\n",
    "\n",
    "### Hierarchical Named Tensors\n",
    "\n",
    "Batch stores **hierarchical named tensors** - collections of tensors whose identifiers form a structured hierarchy. Consider tensors `[t1, t2, t3, t4]` with names `[name1, name2, name3, name4]`, where `name1` and `name2` are under namespace `name0`. The fully qualified name of `t1` is `name0.name1`.\n",
    "\n",
    "### Tree Structure Visualization\n",
    "\n",
    "The structure can be visualized as a tree with:\n",
    "- **Root**: The Batch object itself\n",
    "- **Internal nodes**: Keys (names)\n",
    "- **Leaf nodes**: Values (scalars, arrays, tensors)\n",
    "\n",
    "```mermaid\n",
    "graph TD\n",
    "    root[\"Batch (root)\"]\n",
    "    root --> obs[\"obs\"]\n",
    "    root --> act[\"act\"]\n",
    "    root --> rew[\"rew\"]\n",
    "    obs --> camera[\"camera\"]\n",
    "    obs --> sensory[\"sensory\"]\n",
    "    camera --> cam_data[\"np.array(3,3)\"]\n",
    "    sensory --> sens_data[\"np.array(5,)\"]\n",
    "    act --> act_data[\"np.array(2,)\"]\n",
    "    rew --> rew_data[\"3.66\"]\n",
    "    \n",
    "    style root fill:#e1f5ff\n",
    "    style obs fill:#fff4e1\n",
    "    style act fill:#fff4e1\n",
    "    style rew fill:#fff4e1\n",
    "    style camera fill:#ffe1f5\n",
    "    style sensory fill:#ffe1f5\n",
    "    style cam_data fill:#e8f5e1\n",
    "    style sens_data fill:#e8f5e1\n",
    "    style act_data fill:#e8f5e1\n",
    "    style rew_data fill:#e8f5e1\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example: hierarchical structure\n",
    "data = {\n",
    "    \"action\": np.array([1.0, 2.0, 3.0]),\n",
    "    \"reward\": 3.66,\n",
    "    \"obs\": {\n",
    "        \"camera\": np.zeros((3, 3)),\n",
    "        \"sensory\": np.ones(5),\n",
    "    },\n",
    "}\n",
    "\n",
    "batch = Batch(data)\n",
    "print(batch)\n",
    "print(\"\\nAccessing nested values:\")\n",
    "print(f\"batch.obs.camera.shape = {batch.obs.camera.shape}\")\n",
    "print(f\"batch.obs.sensory = {batch.obs.sensory}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Flow in RL Pipeline\n",
    "\n",
    "Batch facilitates data flow throughout the RL pipeline:\n",
    "\n",
    "```mermaid\n",
    "graph LR\n",
    "    A[Collector] -->|ActBatchProtocol| B[Environment]\n",
    "    B[Environment + Action] -->|RolloutBatchProtocol| C[Replay Buffer]\n",
    "    C -->|RolloutBatchProtocol| D[Policy]\n",
    "    D -->|ActBatchProtocol| A\n",
    "    D -->|BatchWithAdvantages| E[Algorithm/Trainer]\n",
    "    E --> D\n",
    "    \n",
    "    style A fill:#e1f5ff\n",
    "    style B fill:#fff4e1\n",
    "    style C fill:#ffe1f5\n",
    "    style D fill:#e8f5e1\n",
    "    style E fill:#f5e1e1\n",
    "```\n",
    "\n",
    "Each arrow represents a specific `BatchProtocol` that defines what fields are expected at that stage."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Basic Operations\n",
    "\n",
    "### 3.1 Construction\n",
    "\n",
    "Batch objects can be constructed in several ways:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# From keyword arguments\n",
    "batch1 = Batch(a=4, b=[5, 5], c=\"hello\")\n",
    "print(\"From kwargs:\", batch1)\n",
    "\n",
    "# From dictionary\n",
    "batch2 = Batch({\"a\": 4, \"b\": [5, 5], \"c\": \"hello\"})\n",
    "print(\"\\nFrom dict:\", batch2)\n",
    "\n",
    "# From list of dictionaries (automatically stacked)\n",
    "batch3 = Batch([{\"a\": 1, \"b\": 2}, {\"a\": 3, \"b\": 4}])\n",
    "print(\"\\nFrom list of dicts:\", batch3)\n",
    "\n",
    "# Nested batch\n",
    "batch4 = Batch(obs=Batch(x=1, y=2), act=5)\n",
    "print(\"\\nNested:\", batch4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Content Rules\n",
    "\n",
    "Understanding what Batch can store and how it converts data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Keys must be strings\n",
    "batch = Batch()\n",
    "batch.key1 = \"value\"\n",
    "batch.key2 = np.array([1, 2, 3])\n",
    "print(\"Keys:\", list(batch.keys()))\n",
    "\n",
    "# Automatic conversions\n",
    "demo = Batch(\n",
    "    scalar_int=5,  # → np.array(5)\n",
    "    scalar_float=3.14,  # → np.array(3.14)\n",
    "    list_nums=[1, 2, 3],  # → np.array([1, 2, 3])\n",
    "    list_mixed=[1, \"hello\", None],  # → np.array([1, \"hello\", None], dtype=object)\n",
    "    dict_val={\"x\": 1, \"y\": 2},  # → Batch(x=1, y=2)\n",
    ")\n",
    "\n",
    "print(\"\\nAutomatic conversions:\")\n",
    "print(f\"scalar_int type: {type(demo.scalar_int)}, value: {demo.scalar_int}\")\n",
    "print(f\"list_nums type: {type(demo.list_nums)}, dtype: {demo.list_nums.dtype}\")\n",
    "print(f\"list_mixed dtype: {demo.list_mixed.dtype}\")\n",
    "print(f\"dict_val type: {type(demo.dict_val)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Important conversions:**\n",
    "- Lists of numbers → NumPy arrays\n",
    "- Lists with mixed types → Object-dtype arrays\n",
    "- Dictionaries → Batch objects (recursively)\n",
    "- Scalars → NumPy scalars"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 Access Patterns\n",
    "\n",
    "**Important: Understanding Iteration**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = Batch(a=[1, 2, 3], b=[4, 5, 6])\n",
    "\n",
    "# Attribute vs dictionary access (equivalent)\n",
    "print(\"Attribute access:\", batch.a)\n",
    "print(\"Dict access:\", batch[\"a\"])\n",
    "\n",
    "# Getting keys\n",
    "print(\"\\nKeys:\", list(batch.keys()))\n",
    "\n",
    "# Gotcha: Iteration is array like, not over keys\n",
    "print(\"\\nIteration behavior:\")\n",
    "print(\"for x in batch iterates over batch[0], batch[1], ..., NOT keys!\")\n",
    "for i, item in enumerate(batch):\n",
    "    print(f\"batch[{i}] = {item}\")\n",
    "\n",
    "# This is different from dict behavior!\n",
    "regular_dict = {\"a\": [1, 2, 3], \"b\": [4, 5, 6]}\n",
    "print(\"\\nCompare with dict iteration (iterates over keys):\")\n",
    "for key in regular_dict:\n",
    "    print(f\"key = {key}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": ""
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.4 Indexing & Slicing\n",
    "\n",
    "Batch supports NumPy-like indexing and slicing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = Batch(a=np.array([[0.0, 2.0], [1.0, 3.0]]), b=[[5.0, -5.0], [1.0, -2.0]])\n",
    "\n",
    "print(\"Original batch shape:\", batch.shape)\n",
    "print(\"Original batch length:\", len(batch))\n",
    "\n",
    "# Single index\n",
    "print(\"\\nbatch[0]:\")\n",
    "print(batch[0])\n",
    "\n",
    "# Slicing\n",
    "print(\"\\nbatch[:1]:\")\n",
    "print(batch[:1])\n",
    "\n",
    "# Advanced indexing\n",
    "print(\"\\nbatch[[0, 1]]:\")\n",
    "print(batch[[0, 1]])\n",
    "\n",
    "# Multi-dimensional indexing\n",
    "print(\"\\nbatch[:, 0] (first column of all arrays):\")\n",
    "print(batch[:, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Broadcasting and in-place operations\n",
    "batch[:, 1] += 10\n",
    "print(\"After batch[:, 1] += 10:\")\n",
    "print(batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.5 Stack, Concatenate, and Split\n",
    "\n",
    "Combining and splitting batches:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "# Stack: adds a new dimension\nbatch1 = Batch(a=np.array([1, 2]), b=np.array([5, 6]))\nbatch2 = Batch(a=np.array([3, 4]), b=np.array([7, 8]))\n\nstacked = Batch.stack([batch1, batch2])\nprint(\"Stacked:\")\nprint(stacked)\nprint(f\"Shape: {stacked.shape}\")\n\n# Concatenate: extends along existing dimension\nconcatenated = Batch.cat([batch1, batch2])\nprint(\"\\nConcatenated:\")\nprint(concatenated)\nprint(f\"Shape: {concatenated.shape}\")"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split\n",
    "batch = Batch(a=np.arange(10), b=np.arange(10, 20))\n",
    "splits = list(batch.split(size=3, shuffle=False))\n",
    "print(f\"Split into {len(splits)} batches:\")\n",
    "for i, split in enumerate(splits):\n",
    "    print(f\"Split {i}: a={split.a}, length={len(split)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.6 Data Type Conversion\n",
    "\n",
    "Converting between NumPy and PyTorch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create batch with NumPy arrays\n",
    "batch = Batch(a=np.zeros((3, 4)), b=np.ones(5))\n",
    "print(\"Original (NumPy):\")\n",
    "print(f\"batch.a type: {type(batch.a)}\")\n",
    "\n",
    "# Convert to PyTorch (in-place)\n",
    "batch.to_torch_(dtype=torch.float32, device=\"cpu\")\n",
    "print(\"\\nAfter to_torch_():\")\n",
    "print(f\"batch.a type: {type(batch.a)}\")\n",
    "print(f\"batch.a dtype: {batch.a.dtype}\")\n",
    "\n",
    "# Convert back to NumPy (in-place)\n",
    "batch.to_numpy_()\n",
    "print(\"\\nAfter to_numpy_():\")\n",
    "print(f\"batch.a type: {type(batch.a)}\")\n",
    "\n",
    "# Non-in-place versions return a new batch\n",
    "batch_torch = batch.to_torch()\n",
    "print(\"\\nOriginal batch unchanged:\", type(batch.a))\n",
    "print(\"New batch:\", type(batch_torch.a))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Type Safety with Protocols\n",
    "\n",
    "### Why Protocols?\n",
    "\n",
    "Batch needs to be **flexible** (not fixed fields like dataclasses) but we still want **type safety** and **IDE autocompletion**. Protocols provide the best of both worlds:\n",
    "\n",
    "- **Runtime flexibility**: Add any fields dynamically\n",
    "- **Static type checking**: Type checkers (mypy, pyright) verify correct usage\n",
    "- **IDE support**: Autocompletion for expected fields\n",
    "\n",
    "### What is BatchProtocol?\n",
    "\n",
    "A `Protocol` defines an interface without implementation. Think of it as a contract: \"any object with these fields is valid.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a typed batch using cast\n",
    "# This enables IDE autocompletion and type checking\n",
    "\n",
    "# ActBatchProtocol: just needs 'act' field\n",
    "act_batch = cast(ActBatchProtocol, Batch(act=np.array([1, 2, 3])))\n",
    "print(\"ActBatchProtocol:\", act_batch.act)\n",
    "\n",
    "# ObsBatchProtocol: needs 'obs' and 'info' fields\n",
    "obs_batch = cast(\n",
    "    ObsBatchProtocol,\n",
    "    Batch(obs=np.array([[1.0, 2.0], [3.0, 4.0]]), info=np.array([{}, {}], dtype=object)),\n",
    ")\n",
    "print(\"\\nObsBatchProtocol:\", obs_batch.obs)\n",
    "\n",
    "# RolloutBatchProtocol: needs obs, obs_next, act, rew, terminated, truncated\n",
    "rollout_batch = cast(\n",
    "    RolloutBatchProtocol,\n",
    "    Batch(\n",
    "        obs=np.array([[1.0, 2.0], [3.0, 4.0]]),\n",
    "        obs_next=np.array([[2.0, 3.0], [4.0, 5.0]]),\n",
    "        act=np.array([0, 1]),\n",
    "        rew=np.array([1.0, 2.0]),\n",
    "        terminated=np.array([False, True]),\n",
    "        truncated=np.array([False, False]),\n",
    "        info=np.array([{}, {}], dtype=object),\n",
    "    ),\n",
    ")\n",
    "print(\"\\nRolloutBatchProtocol reward:\", rollout_batch.rew)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Protocol Hierarchy\n",
    "\n",
    "Tianshou defines a hierarchy of protocols for different use cases:\n",
    "\n",
    "```mermaid\n",
    "graph TD\n",
    "    BP[BatchProtocol<br/>Base protocol] --> OBP[ObsBatchProtocol<br/>obs, info]\n",
    "    BP --> ABP[ActBatchProtocol<br/>act]\n",
    "    ABP --> ASBP[ActStateBatchProtocol<br/>act, state]\n",
    "    OBP --> RBP[RolloutBatchProtocol<br/>+obs_next, act, rew,<br/>terminated, truncated]\n",
    "    RBP --> BWRP[BatchWithReturnsProtocol<br/>+returns]\n",
    "    BWRP --> BWAP[BatchWithAdvantagesProtocol<br/>+adv, v_s]\n",
    "    ASBP --> MOBP[ModelOutputBatchProtocol<br/>+logits]\n",
    "    MOBP --> DBP[DistBatchProtocol<br/>+dist]\n",
    "    DBP --> DLPBP[DistLogProbBatchProtocol<br/>+log_prob]\n",
    "    BWAP --> LOPBP[LogpOldProtocol<br/>+logp_old]\n",
    "    \n",
    "    style BP fill:#e1f5ff\n",
    "    style OBP fill:#fff4e1\n",
    "    style ABP fill:#fff4e1\n",
    "    style RBP fill:#ffe1f5\n",
    "    style BWRP fill:#e8f5e1\n",
    "    style BWAP fill:#e8f5e1\n",
    "    style DBP fill:#f5e1e1\n",
    "    style LOPBP fill:#e1e1f5\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using Protocols in Functions\n",
    "\n",
    "Protocols enable type-safe function signatures:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_observations(batch: ObsBatchProtocol) -> np.ndarray:\n",
    "    \"\"\"Function that expects observations.\n",
    "\n",
    "    IDE will autocomplete batch.obs and batch.info!\n",
    "    Type checker will verify these fields exist.\n",
    "    \"\"\"\n",
    "    # IDE knows batch.obs exists\n",
    "    return batch.obs if isinstance(batch.obs, np.ndarray) else np.array(batch.obs)\n",
    "\n",
    "\n",
    "def compute_advantage(batch: RolloutBatchProtocol) -> np.ndarray:\n",
    "    \"\"\"Function that expects rollout data.\n",
    "\n",
    "    IDE will autocomplete batch.rew, batch.obs_next, etc.\n",
    "    \"\"\"\n",
    "    # Simplified advantage computation\n",
    "    return batch.rew  # IDE knows this exists\n",
    "\n",
    "\n",
    "# Example usage\n",
    "obs_data = Batch(obs=np.array([1, 2, 3]), info=np.array([{}], dtype=object))\n",
    "result = process_observations(obs_data)\n",
    "print(\"Processed obs:\", result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Key Protocol Types:**\n",
    "\n",
    "- `ActBatchProtocol`: Just actions (for simple policies)\n",
    "- `ObsBatchProtocol`: Observations and info\n",
    "- `RolloutBatchProtocol`: Complete transitions (obs, act, rew, done, obs_next)\n",
    "- `BatchWithReturnsProtocol`: Rollouts + computed returns\n",
    "- `BatchWithAdvantagesProtocol`: Returns + advantages and values\n",
    "- `DistBatchProtocol`: Contains distribution objects\n",
    "- `LogpOldProtocol`: For importance sampling (PPO, etc.)\n",
    "\n",
    "See `tianshou/data/types.py` for the complete list!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Distribution Slicing\n",
    "\n",
    "### Why Special Handling?\n",
    "\n",
    "PyTorch `Distribution` objects need special slicing because they're not simple arrays. When you slice `batch[0:2]`, Tianshou needs to slice the underlying distribution parameters correctly.\n",
    "\n",
    "### Supported Distributions\n",
    "\n",
    "Tianshou supports slicing for:\n",
    "- `Categorical`: Discrete distributions\n",
    "- `Normal`: Continuous Gaussian distributions\n",
    "- `Independent`: Wraps other distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Categorical distribution\n",
    "probs = torch.tensor([[0.3, 0.7], [0.4, 0.6], [0.5, 0.5]])\n",
    "dist = Categorical(probs=probs)\n",
    "batch = Batch(dist=dist, values=np.array([1, 2, 3]))\n",
    "\n",
    "print(\"Original batch length:\", len(batch))\n",
    "print(\"Original dist probs shape:\", batch.dist.probs.shape)\n",
    "\n",
    "# Slicing automatically handles the distribution\n",
    "sliced = batch[0:2]\n",
    "print(\"\\nSliced batch length:\", len(sliced))\n",
    "print(\"Sliced dist probs shape:\", sliced.dist.probs.shape)\n",
    "print(\"Sliced values:\", sliced.values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normal distribution\n",
    "loc = torch.tensor([0.0, 1.0, 2.0])\n",
    "scale = torch.tensor([1.0, 1.0, 1.0])\n",
    "normal_dist = Normal(loc=loc, scale=scale)\n",
    "batch_normal = Batch(dist=normal_dist, actions=np.array([0.5, 1.5, 2.5]))\n",
    "\n",
    "print(\"Normal distribution batch:\")\n",
    "print(f\"Original mean: {batch_normal.dist.mean}\")\n",
    "\n",
    "# Index a single element\n",
    "single = batch_normal[1]\n",
    "print(f\"\\nSingle element mean: {single.dist.mean}\")\n",
    "print(f\"Single element action: {single.actions}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Converting to At Least 2D\n",
    "\n",
    "Sometimes you need to ensure distributions have a batch dimension:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tianshou.data.batch import dist_to_atleast_2d\n",
    "\n",
    "# Scalar distribution (no batch dimension)\n",
    "scalar_dist = Categorical(probs=torch.tensor([0.3, 0.7]))\n",
    "print(\"Scalar dist batch_shape:\", scalar_dist.batch_shape)\n",
    "\n",
    "# Convert to have batch dimension\n",
    "batched_dist = dist_to_atleast_2d(scalar_dist)\n",
    "print(\"Batched dist batch_shape:\", batched_dist.batch_shape)\n",
    "\n",
    "# For entire batch\n",
    "scalar_batch = Batch(a=1, b=2, dist=Categorical(probs=torch.ones(3)))\n",
    "print(\"\\nBefore to_at_least_2d:\", scalar_batch.dist.batch_shape)\n",
    "\n",
    "batch_2d = scalar_batch.to_at_least_2d()\n",
    "print(\"After to_at_least_2d:\", batch_2d.dist.batch_shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use Cases\n",
    "\n",
    "Distribution slicing is used in:\n",
    "- **Policy sampling**: When policies output distributions, slicing batches preserves distribution structure\n",
    "- **Replay buffer sampling**: Distributions are stored and retrieved correctly\n",
    "- **Advantage computation**: Computing log probabilities on subsets of data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Advanced Topics\n",
    "\n",
    "### 6.1 Key Reservation\n",
    "\n",
    "Sometimes you know what keys you'll need but don't have values yet. Reserve keys using empty `Batch()` objects:\n",
    "\n",
    "```mermaid\n",
    "graph TD\n",
    "    root[\"Batch\"]\n",
    "    root --> a[\"key1: np.array([1,2,3])\"]\n",
    "    root --> b[\"key2: Batch() (reserved)\"]\n",
    "    root --> c[\"key3\"]\n",
    "    c --> c1[\"subkey1: Batch() (reserved)\"]\n",
    "    c --> c2[\"subkey2: np.array([4,5])\"]\n",
    "    \n",
    "    style root fill:#e1f5ff\n",
    "    style a fill:#e8f5e1\n",
    "    style b fill:#ffcccc\n",
    "    style c fill:#fff4e1\n",
    "    style c1 fill:#ffcccc\n",
    "    style c2 fill:#e8f5e1\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reserving keys\n",
    "batch = Batch(\n",
    "    known_field=np.array([1, 2]),\n",
    "    future_field=Batch(),  # Reserved for later\n",
    ")\n",
    "print(\"Batch with reserved key:\")\n",
    "print(batch)\n",
    "\n",
    "# Later, assign actual data\n",
    "batch.future_field = np.array([3, 4])\n",
    "print(\"\\nAfter assignment:\")\n",
    "print(batch)\n",
    "\n",
    "# Nested reservation\n",
    "batch2 = Batch(\n",
    "    obs=Batch(\n",
    "        camera=Batch(),  # Reserved\n",
    "        lidar=np.zeros(10),\n",
    "    )\n",
    ")\n",
    "print(\"\\nNested reservation:\")\n",
    "print(batch2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.2 Length and Shape Semantics\n",
    "\n",
    "Understanding when `len()` works and what `shape` means:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normal case: all tensors same length\n",
    "batch1 = Batch(a=[1, 2, 3], b=np.array([4, 5, 6]))\n",
    "print(\"Normal batch:\")\n",
    "print(f\"len(batch1) = {len(batch1)}\")\n",
    "print(f\"batch1.shape = {batch1.shape}\")\n",
    "\n",
    "# Scalars have no length\n",
    "batch2 = Batch(a=5, b=10)\n",
    "print(\"\\nScalar batch:\")\n",
    "print(f\"batch2.shape = {batch2.shape}\")\n",
    "try:\n",
    "    print(f\"len(batch2) = {len(batch2)}\")\n",
    "except TypeError as e:\n",
    "    print(f\"len(batch2) raises TypeError: {e}\")\n",
    "\n",
    "# Mixed lengths: returns minimum\n",
    "batch3 = Batch(a=[1, 2], b=[3, 4, 5])\n",
    "print(\"\\nMixed length batch:\")\n",
    "print(f\"len(batch3) = {len(batch3)} (minimum of 2 and 3)\")\n",
    "\n",
    "# Reserved keys are ignored\n",
    "batch4 = Batch(a=[1, 2, 3], reserved=Batch())\n",
    "print(\"\\nBatch with reserved key:\")\n",
    "print(f\"len(batch4) = {len(batch4)} (reserved key ignored)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.3 Empty Batches\n",
    "\n",
    "Understanding different meanings of \"empty\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. No keys at all\n",
    "empty1 = Batch()\n",
    "print(\"No keys:\")\n",
    "print(f\"len(empty1.get_keys()) = {len(list(empty1.get_keys()))}\")\n",
    "print(f\"len(empty1) = {len(empty1)}\")\n",
    "\n",
    "# 2. Has keys but they're all reserved\n",
    "empty2 = Batch(a=Batch(), b=Batch())\n",
    "print(\"\\nReserved keys only:\")\n",
    "print(f\"len(empty2.get_keys()) = {len(list(empty2.get_keys()))}\")\n",
    "print(f\"len(empty2) = {len(empty2)}\")\n",
    "\n",
    "# 3. Has data but length is 0\n",
    "empty3 = Batch(a=np.array([]), b=np.array([]))\n",
    "print(\"\\nZero-length arrays:\")\n",
    "print(f\"len(empty3.get_keys()) = {len(list(empty3.get_keys()))}\")\n",
    "print(f\"len(empty3) = {len(empty3)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Checking emptiness:**\n",
    "- `len(batch.get_keys()) == 0`: No keys (completely empty)\n",
    "- `len(batch) == 0`: No data elements (may have reserved keys)\n",
    "\n",
    "**The `.empty()` and `.empty_()` methods:**\n",
    "These reset values to zeros/None, different from checking emptiness:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = Batch(a=[1, 2, 3], b=[\"x\", \"y\", \"z\"])\n",
    "print(\"Original:\", batch)\n",
    "\n",
    "# Empty specific index\n",
    "batch[0] = Batch.empty(batch[0])\n",
    "print(\"\\nAfter emptying index 0:\")\n",
    "print(batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.4 Heterogeneous Aggregation\n",
    "\n",
    "Stacking/concatenating batches with different keys:\n",
    "\n",
    "```mermaid\n",
    "graph LR\n",
    "    A[\"Batch(a=[1,2], c=5)\"] --> C[\"Batch.stack\"]\n",
    "    B[\"Batch(b=[3,4], c=6)\"] --> C\n",
    "    C --> D[\"Batch(a=[[1,2],[0,0]],<br/>b=[[0,0],[3,4]],<br/>c=[5,6])\"]\n",
    "    \n",
    "    style A fill:#e1f5ff\n",
    "    style B fill:#fff4e1\n",
    "    style C fill:#ffe1f5\n",
    "    style D fill:#e8f5e1\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stack with different keys (missing keys padded with zeros)\n",
    "batch_a = Batch(a=np.ones((2, 3)), shared=np.array([1, 2]))\n",
    "batch_b = Batch(b=np.zeros((2, 4)), shared=np.array([3, 4]))\n",
    "\n",
    "stacked = Batch.stack([batch_a, batch_b])\n",
    "print(\"Stacked batch:\")\n",
    "print(f\"a.shape = {stacked.a.shape} (padded with zeros for batch_b)\")\n",
    "print(f\"b.shape = {stacked.b.shape} (padded with zeros for batch_a)\")\n",
    "print(f\"shared.shape = {stacked.shared.shape} (in both batches)\")\n",
    "print(stacked)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.5 Missing Values\n",
    "\n",
    "Handling `None` and `NaN` values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Batch with missing values\n",
    "batch = Batch(a=[1, 2, None, 4], b=[5.0, np.nan, 7.0, 8.0], c=[[1, 2], [3, 4], [5, 6], [7, 8]])\n",
    "\n",
    "# Check for nulls\n",
    "print(\"Has null?\", batch.hasnull())\n",
    "\n",
    "# Get null mask\n",
    "null_mask = batch.isnull()\n",
    "print(\"\\nNull mask:\")\n",
    "print(f\"a: {null_mask.a}\")\n",
    "print(f\"b: {null_mask.b}\")\n",
    "\n",
    "# Drop rows with any null\n",
    "clean_batch = batch.dropnull()\n",
    "print(\"\\nAfter dropnull() (keeps rows 0 and 3):\")\n",
    "print(f\"Length: {len(clean_batch)}\")\n",
    "print(f\"a: {clean_batch.a}\")\n",
    "print(f\"b: {clean_batch.b}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.6 Value Transformations\n",
    "\n",
    "Applying functions to all values recursively:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = Batch(a=np.array([1, 2, 3]), nested=Batch(b=np.array([4.0, 5.0]), c=np.array([6, 7, 8])))\n",
    "\n",
    "# Apply transformation (returns new batch)\n",
    "doubled = batch.apply_values_transform(lambda x: x * 2)\n",
    "print(\"Original batch a:\", batch.a)\n",
    "print(\"Doubled batch a:\", doubled.a)\n",
    "print(\"Doubled nested.b:\", doubled.nested.b)\n",
    "\n",
    "# In-place transformation\n",
    "batch.apply_values_transform(lambda x: x + 10, inplace=True)\n",
    "print(\"\\nAfter in-place +10:\")\n",
    "print(\"a:\", batch.a)\n",
    "print(\"nested.b:\", batch.nested.b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Surprising Behaviors & Gotchas\n",
    "\n",
    "### Iteration Does NOT Iterate Over Keys!\n",
    "\n",
    "**This is the most common source of confusion:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = Batch(a=[1, 2, 3], b=[4, 5, 6])\n",
    "\n",
    "print(\"WRONG: This doesn't iterate over keys!\")\n",
    "for item in batch:\n",
    "    print(f\"item = {item}\")  # Prints batch[0], batch[1], batch[2]\n",
    "\n",
    "print(\"\\nCORRECT: To iterate over keys:\")\n",
    "for key in batch.keys():\n",
    "    print(f\"key = {key}\")\n",
    "\n",
    "print(\"\\nCORRECT: To iterate over key-value pairs:\")\n",
    "for key, value in batch.items():\n",
    "    print(f\"{key} = {value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Automatic Type Conversions\n",
    "\n",
    "Be aware of these automatic conversions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lists become arrays\n",
    "batch = Batch(a=[1, 2, 3])\n",
    "print(\"List → array:\", type(batch.a), batch.a.dtype)\n",
    "\n",
    "# Dicts become Batch\n",
    "batch = Batch(a={\"x\": 1, \"y\": 2})\n",
    "print(\"Dict → Batch:\", type(batch.a))\n",
    "\n",
    "# Scalars become numpy scalars\n",
    "batch = Batch(a=5)\n",
    "print(\"Scalar → np.ndarray:\", type(batch.a), batch.a)\n",
    "\n",
    "# Mixed types → object dtype\n",
    "batch = Batch(a=[1, \"hello\", None])\n",
    "print(\"Mixed → object:\", batch.a.dtype, batch.a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Length Edge Cases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Scalars have no length\n",
    "batch_scalar = Batch(a=5, b=10)\n",
    "try:\n",
    "    len(batch_scalar)\n",
    "except TypeError as e:\n",
    "    print(f\"Scalar batch: {e}\")\n",
    "\n",
    "# 2. Empty nested batches ignored in len()\n",
    "batch_empty_nested = Batch(a=[1, 2, 3], b=Batch())\n",
    "print(f\"\\nWith empty nested: len = {len(batch_empty_nested)} (ignores b)\")\n",
    "\n",
    "# 3. Different lengths: returns minimum\n",
    "batch_different = Batch(a=[1, 2], b=[1, 2, 3, 4])\n",
    "print(f\"Different lengths: len = {len(batch_different)} (minimum)\")\n",
    "\n",
    "# 4. None values don't affect length\n",
    "batch_none = Batch(a=[1, 2, 3], b=None)\n",
    "print(f\"With None: len = {len(batch_none)} (None ignored)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### String Keys Only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Integer keys not allowed\n",
    "try:\n",
    "    batch = Batch({1: \"value\", 2: \"other\"})\n",
    "except AssertionError as e:\n",
    "    print(\"Integer keys not allowed:\", e)\n",
    "\n",
    "# String keys work\n",
    "batch = Batch({\"key1\": \"value\", \"key2\": \"other\"})\n",
    "print(\"\\nString keys work:\", list(batch.keys()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cat vs Stack Behavior\n",
    "\n",
    "Recent changes have made concatenation stricter about structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stack pads missing keys with zeros\n",
    "b1 = Batch(a=[1, 2])\n",
    "b2 = Batch(b=[3, 4])\n",
    "stacked = Batch.stack([b1, b2])\n",
    "print(\"Stack (different keys):\")\n",
    "print(f\"  a: {stacked.a}  (b2.a padded with 0)\")\n",
    "print(f\"  b: {stacked.b}  (b1.b padded with 0)\")\n",
    "\n",
    "# Cat requires same structure now\n",
    "b3 = Batch(a=[1, 2], b=[3, 4])\n",
    "b4 = Batch(a=[5, 6], b=[7, 8])\n",
    "concatenated = Batch.cat([b3, b4])\n",
    "print(\"\\nCat (same keys):\")\n",
    "print(f\"  a: {concatenated.a}\")\n",
    "print(f\"  b: {concatenated.b}\")\n",
    "\n",
    "# Cat with different structures raises error\n",
    "try:\n",
    "    Batch.cat([b1, b2])  # Different keys!\n",
    "except ValueError:\n",
    "    print(\"\\nCat with different keys: ValueError raised\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Best Practices\n",
    "\n",
    "### When to Use Batch\n",
    "\n",
    "**Good use cases:**\n",
    "- Collecting environment data (transitions, episodes)\n",
    "- Storing replay buffer data\n",
    "- Passing data between components (collector → buffer → policy)\n",
    "- Handling heterogeneous observations (dict spaces)\n",
    "\n",
    "**Consider alternatives:**\n",
    "- Simple scalar tracking (use regular variables)\n",
    "- Pure tensor operations (use PyTorch tensors directly)\n",
    "- Deeply nested arbitrary structures (use dataclasses)\n",
    "\n",
    "### Structuring Your Batches\n",
    "\n",
    "**Use protocols for type safety:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Good: Use protocols for clear interfaces\n",
    "def train_step(batch: RolloutBatchProtocol) -> float:\n",
    "    \"\"\"IDE knows what fields exist.\"\"\"\n",
    "    loss = ((batch.rew - 0.5) ** 2).mean()  # Type-safe\n",
    "    return float(loss)\n",
    "\n",
    "\n",
    "# Create properly typed batch\n",
    "train_batch = cast(\n",
    "    RolloutBatchProtocol,\n",
    "    Batch(\n",
    "        obs=np.random.randn(10, 4),\n",
    "        obs_next=np.random.randn(10, 4),\n",
    "        act=np.random.randint(0, 2, 10),\n",
    "        rew=np.random.randn(10),\n",
    "        terminated=np.zeros(10, dtype=bool),\n",
    "        truncated=np.zeros(10, dtype=bool),\n",
    "        info=np.array([{}] * 10, dtype=object),\n",
    "    ),\n",
    ")\n",
    "\n",
    "loss = train_step(train_batch)\n",
    "print(f\"Loss: {loss:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Consistent key naming:**\n",
    "- Follow Tianshou conventions: `obs`, `act`, `rew`, `terminated`, `truncated`\n",
    "- Use descriptive names: `camera_obs` not `co`\n",
    "- Avoid name collisions with Batch methods: don't use `keys`, `items`, `get`, etc.\n",
    "\n",
    "**When to nest vs flatten:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Good: Nest related data\n",
    "batch_nested = Batch(\n",
    "    obs=Batch(\n",
    "        camera=np.zeros((32, 64, 64, 3)), lidar=np.zeros((32, 100)), position=np.zeros((32, 3))\n",
    "    ),\n",
    "    act=np.zeros(32),\n",
    ")\n",
    "print(\"Nested structure for related obs:\")\n",
    "print(f\"  Access: batch.obs.camera.shape = {batch_nested.obs.camera.shape}\")\n",
    "\n",
    "# Less good: Flat structure loses semantic grouping\n",
    "batch_flat = Batch(\n",
    "    camera=np.zeros((32, 64, 64, 3)),\n",
    "    lidar=np.zeros((32, 100)),\n",
    "    position=np.zeros((32, 3)),\n",
    "    act=np.zeros(32),\n",
    ")\n",
    "print(\"\\nFlat structure (works but less clear):\")\n",
    "print(f\"  Access: batch.camera.shape = {batch_flat.camera.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Performance Tips\n",
    "\n",
    "**Use in-place operations:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "batch = Batch(a=np.random.randn(1000, 100))\n",
    "\n",
    "# Creates copy\n",
    "start = time.time()\n",
    "for _ in range(100):\n",
    "    _ = batch.to_torch()\n",
    "time_copy = time.time() - start\n",
    "\n",
    "# In-place (faster)\n",
    "start = time.time()\n",
    "for _ in range(100):\n",
    "    batch.to_torch_()\n",
    "    batch.to_numpy_()\n",
    "time_inplace = time.time() - start\n",
    "\n",
    "print(f\"Copy: {time_copy:.4f}s\")\n",
    "print(f\"In-place: {time_inplace:.4f}s\")\n",
    "print(f\"Speedup: {time_copy / time_inplace:.1f}x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Be mindful of copies:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.array([1, 2, 3])\n",
    "\n",
    "# Default: creates reference (be careful!)\n",
    "batch1 = Batch(a=arr)\n",
    "batch1.a[0] = 999\n",
    "print(f\"Original array modified: {arr}\")  # Changed!\n",
    "\n",
    "# Explicit copy when needed\n",
    "arr = np.array([1, 2, 3])\n",
    "batch2 = Batch(a=arr, copy=True)\n",
    "batch2.a[0] = 999\n",
    "print(f\"Original array preserved: {arr}\")  # Unchanged"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Avoid unnecessary conversions:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inefficient: multiple conversions\n",
    "batch = Batch(a=np.random.randn(100, 10))\n",
    "batch.to_torch_()\n",
    "batch.to_numpy_()  # Unnecessary if we just need NumPy\n",
    "\n",
    "# Efficient: convert once, use many times\n",
    "batch = Batch(a=np.random.randn(100, 10))\n",
    "batch.to_torch_()  # Convert once\n",
    "# ... do torch operations ...\n",
    "# Keep as torch if that's what you need!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Common Patterns\n",
    "\n",
    "**Pattern 1: Building batches incrementally**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Collect data from multiple steps\n",
    "step_data = []\n",
    "for i in range(5):\n",
    "    step_data.append({\"obs\": np.random.randn(4), \"act\": i, \"rew\": np.random.randn()})\n",
    "\n",
    "# Convert to batch (automatically stacks)\n",
    "episode_batch = Batch(step_data)\n",
    "print(\"Episode batch shape:\", episode_batch.shape)\n",
    "print(\"obs shape:\", episode_batch.obs.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Pattern 2: Slicing for mini-batches**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Large batch\n",
    "large_batch = Batch(obs=np.random.randn(100, 4), act=np.random.randint(0, 2, 100))\n",
    "\n",
    "# Split into mini-batches\n",
    "batch_size = 32\n",
    "for mini_batch in large_batch.split(batch_size, shuffle=True):\n",
    "    print(f\"Mini-batch size: {len(mini_batch)}\")\n",
    "    # Train on mini_batch...\n",
    "    break  # Just show one iteration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Pattern 3: Extending batches**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start with some data\n",
    "batch = Batch(obs=np.array([[1, 2], [3, 4]]), act=np.array([0, 1]))\n",
    "print(\"Initial:\", len(batch))\n",
    "\n",
    "# Add more data\n",
    "new_data = Batch(obs=np.array([[5, 6]]), act=np.array([1]))\n",
    "batch.cat_(new_data)\n",
    "print(\"After cat_:\", len(batch))\n",
    "print(\"obs:\", batch.obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Summary\n",
    "\n",
    "### Key Takeaways\n",
    "\n",
    "1. **Batch = Dict + Array**: Combines key-value storage with array operations\n",
    "2. **Hierarchical Structure**: Perfect for complex RL data (nested observations, etc.)\n",
    "3. **Type Safety via Protocols**: Use `BatchProtocol` subclasses for IDE support and type checking\n",
    "4. **Special RL Features**: Distribution slicing, heterogeneous aggregation, missing value handling\n",
    "5. **Remember**: Iteration is over indices, NOT keys!\n",
    "\n",
    "### Quick Reference\n",
    "\n",
    "| Operation | Code | Notes |\n",
    "|-----------|------|-------|\n",
    "| Create | `Batch(a=1, b=[2, 3])` | Auto-converts types |\n",
    "| Access | `batch.a` or `batch[\"a\"]` | Equivalent |\n",
    "| Index | `batch[0]`, `batch[:10]` | Returns sliced Batch |\n",
    "| Iterate indices | `for item in batch:` | Yields batch[0], batch[1], ... |\n",
    "| Iterate keys | `for k in batch.keys():` | Like dict |\n",
    "| Stack | `Batch.stack([b1, b2])` | Adds dimension |\n",
    "| Concatenate | `Batch.cat([b1, b2])` | Extends dimension |\n",
    "| Split | `batch.split(size=10)` | Returns iterator |\n",
    "| To PyTorch | `batch.to_torch_()` | In-place |\n",
    "| To NumPy | `batch.to_numpy_()` | In-place |\n",
    "| Transform | `batch.apply_values_transform(fn)` | Recursive |\n",
    "\n",
    "### Next Steps\n",
    "\n",
    "- **Collector Deep Dive**: See how Batch flows through data collection\n",
    "- **Buffer Deep Dive**: Understand how Batch is stored and sampled\n",
    "- **Policy Guide**: Learn how policies work with BatchProtocol\n",
    "- **API Reference**: Full details at [Batch API documentation](https://tianshou.org/en/stable/api/tianshou.data.html#tianshou.data.Batch)\n",
    "\n",
    "### Questions?\n",
    "\n",
    "- Check the [Tianshou GitHub discussions](https://github.com/thu-ml/tianshou/discussions)\n",
    "- Review [issue tracker](https://github.com/thu-ml/tianshou/issues) for known gotchas\n",
    "- Read the [source code](https://github.com/thu-ml/tianshou/blob/master/tianshou/data/batch.py) - it's well-documented!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Appendix: Serialization & Advanced Topics\n",
    "\n",
    "### Pickle Support"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Batch objects are picklable\n",
    "original = Batch(obs=Batch(a=0.0, c=torch.Tensor([1.0, 2.0])), np=np.zeros([3, 4]))\n",
    "\n",
    "# Serialize and deserialize\n",
    "serialized = pickle.dumps(original)\n",
    "restored = pickle.loads(serialized)\n",
    "\n",
    "print(\"Original obs.a:\", original.obs.a)\n",
    "print(\"Restored obs.a:\", restored.obs.a)\n",
    "print(\"Equal:\", original == restored)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Advanced Indexing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Multi-dimensional data\n",
    "batch = Batch(a=np.random.randn(5, 3, 2))\n",
    "print(\"Original shape:\", batch.a.shape)\n",
    "\n",
    "# Various indexing operations\n",
    "print(\"batch[0].a.shape:\", batch[0].a.shape)\n",
    "print(\"batch[:, 0].a.shape:\", batch[:, 0].a.shape)\n",
    "print(\"batch[[0, 2, 4]].a.shape:\", batch[[0, 2, 4]].a.shape)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
